import os, sys
import numpy as np
import scipy as sp
import matplotlib.pylab as plt

def getWeightsTwoFrames(OS_wave_dir, 
                        nsample_per_pixel=100, 
                        file_num=90, 
                        n_ref=37, 
                        n_sig=37, 
                        mask=None, 
                        out_fname=None):
    """
    Optimize the DCDS weights from two frames A and B.
    """

    if mask is None:
        print('mask is None, quit.')
        sys.exit(0)
    
    if np.sum(mask) != 0:
        print('sum(mask) dose not equal to zero, quit now!')
        sys.exit(0)

    pixel_num = None

    # 1) load waveform data:
    d = None
    dd = None
    for i in range(file_num):
        wfile_A = os.path.join(OS_wave_dir,str(i+1)+'_3.txt')
        wfile_B = os.path.join(OS_wave_dir,str(i+1)+'_4.txt')
        wfile_C = os.path.join(OS_wave_dir,str(i+1)+'_1.txt')
        wfile_D = os.path.join(OS_wave_dir,str(i+1)+'_2.txt')
        
        waveform_A = np.loadtxt(wfile_A)
        waveform_B = np.loadtxt(wfile_B)
        waveform_C = np.loadtxt(wfile_C)
        waveform_D = np.loadtxt(wfile_D)
        
        pixel_num = len(waveform_A)//nsample_per_pixel
        waveform_A = waveform_A.reshape(pixel_num,nsample_per_pixel)
        waveform_B = waveform_B.reshape(pixel_num,nsample_per_pixel)
        waveform_C = waveform_C.reshape(pixel_num,nsample_per_pixel)
        waveform_D = waveform_D.reshape(pixel_num,nsample_per_pixel)
        
        if i==0:
            d = waveform_A - waveform_B
            dd = waveform_C - waveform_D
        else:
            d = np.vstack((d, waveform_A - waveform_B))
            dd = np.vstack((dd, waveform_C - waveform_D))
    
    idx = mask != 0
    d = d[:,idx]
    dd = dd[:,idx]
    
    # 2) prepare matrix N
    N = np.identity(file_num*pixel_num)
    N = N - 1./(file_num*pixel_num)
    
    # 3) prepare matrix A
    A = np.matmul(N,d)
    AtA = np.matmul(A.transpose(), A)
    eA = np.zeros((AtA.shape[0]+2, AtA.shape[1]+2))
    eA[0:AtA.shape[0],0:AtA.shape[1]] = AtA

    i = n_ref + n_sig + 0
    for j in range(n_ref):
        eA[i,j] = 1
        eA[j,i] = 1
    
    i = n_ref + n_sig + 1
    for j in range(n_ref,n_ref+n_sig):
        eA[i,j] = 1
        eA[j,i] = 1

    # 4) prepare vector b
    b = np.zeros(n_ref+n_sig+2)
    b[n_ref+n_sig+0] = 1
    b[n_ref+n_sig+1] = -1
    
    w = np.linalg.solve(eA, b)

    if out_fname is not None:
        wout = []
        w_cnt = 0
        for i in range(nsample_per_pixel):
            if mask[i] != 0:
                wout.append(w[w_cnt])
                w_cnt += 1
            else:
                wout.append(0)
        
        wout = np.array(wout)
        np.savetxt(out_fname, wout, fmt='%10.8f')

    # 5) compare chi2
    # 需要注意将14位“转成”16位，与bin2fits保持一致
    pv_old = ( d[:,0:n_ref].sum(axis=1) - d[:,n_ref:2*n_ref].sum(axis=1) ) / n_ref
    pv_new = np.matmul(d,w[:-2])
    pv_old *= 4
    pv_new *= 4
    std_old = np.std(pv_old) / np.sqrt(2)
    std_new = np.std(pv_new) / np.sqrt(2)

    pv_old2 = ( dd[:,0:n_ref].sum(axis=1) - dd[:,n_ref:2*n_ref].sum(axis=1) ) / n_ref
    pv_new2 = np.matmul(dd,w[:-2])
    pv_old2 *= 4
    pv_new2 *= 4
    std_old2 = np.std(pv_old2) / np.sqrt(2)
    std_new2 = np.std(pv_new2) / np.sqrt(2)

    return std_old, std_new, std_old2, std_new2


